import numpy as np
import torch

from lowpass_filter import lowpass_filter


def wdm_multiplex(Signal, input_signal):
    Sampling_num = Signal['Sampling_num']
    df_matrix = Signal['df_matrix']
    simu_time = Signal['simu_time']
    WDM_channel = Signal['WDM_channel']
    channel_spacing_Hz = Signal['channel_spacing_Hz']

    WDM_sigout = torch.zeros(Sampling_num, dtype=torch.complex64)

    for k in range(WDM_channel):
        phase_array = 2 * np.pi * df_matrix[k] * simu_time

        # Ensure phase_array is a numpy array
        if isinstance(phase_array, torch.Tensor):
            phase_array = phase_array.numpy()

        phase_array_np = np.exp(1j * phase_array).astype(np.complex64)

        h = lowpass_filter(Signal, channel_spacing_Hz)

        # Convert input_signal to complex if it's not already
        if np.iscomplexobj(input_signal):
            input_signal_complex = input_signal
        else:
            # Convert input_signal to complex
            input_signal_complex = input_signal['real'] + 1j * input_signal['imag']

        E_fft_X = np.fft.fftshift(np.fft.fft(input_signal_complex[:, k]))
        E_fft_X *= h
        E_X = np.fft.ifft(np.fft.ifftshift(E_fft_X))

        WDM_sig_X = E_X * phase_array_np
        WDM_sigout += WDM_sig_X

    return WDM_sigout
